import time

from backend.v3_5 import *
import torch
from torch import nn
from torch.nn import functional as F
from collections import OrderedDict
BatchNorm2d = nn.BatchNorm2d

def upsample(x, scale_factor=2, mode="bilinear", align_corners=False):
    """Upsample input tensor by a factor of 2
    """
    return F.interpolate(x, scale_factor=scale_factor, mode=mode, align_corners=align_corners)

def up_cat(x_low, x):
    x_low = F.interpolate(x_low, scale_factor=2,mode='bilinear')
    return torch.cat([x_low, x],dim=1)

def up_add(x_low, x):
    x_low = F.interpolate(x_low, scale_factor=2,mode='bilinear')
    return x_low+x


'''

def up_cat(x_low, x):
    h, w = x.size(2), x.size(3)
    x_low = F.interpolate(x_low, size=(h,w),mode='bilinear')
    return torch.cat([x_low, x],dim=1)
'''

class ConvNormAct(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=(0, 0), dilation=(1, 1), bias=False):
        super(ConvNormAct, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=padding, stride=stride,
                      bias=bias,dilation=dilation),
            BatchNorm2d(out_channels, momentum=0.1),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        x = self.conv(x)
        return x


class SegUNet(nn.Module):
    def __init__(self,num_class,   load_pretrained=False, is_edge=False):
        super(SegUNet, self).__init__()
        self.num_class = num_class
        self.load_pretrained = load_pretrained
        self.backend = resnet_xag_18(is_edge)
        self.inter_channel = 128
        self.is_edge = is_edge
        #
        #self.fea_dim = [80, 160, 320, 640]
        #self.fea_dim = [32, 64, 128, 256]
        #self.fea_dim = [48, 96, 192, 384]
        self.fea_dim = [64, 128, 256, 512]
        self.down_1_2 =  nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
        )
        self.down_1_3 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=2, stride=2),
        )
        self.down_1_4 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=2, stride=2),
        )

        self.down_2_3 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
        )
        self.down_2_4 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=2, stride=2),
        )
        self.down_3_4 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
        )



        self.up_2_1 = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            ConvNormAct(self.fea_dim[1],self.fea_dim[0],kernel_size=1)
        )
        self.up_3_1 = nn.Sequential(
            ConvNormAct(self.fea_dim[2], self.fea_dim[0], kernel_size=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )
        self.up_4_1 = nn.Sequential(
            ConvNormAct(self.fea_dim[3], self.fea_dim[0], kernel_size=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),

        )

        self.up_3_2 = nn.Sequential(
            ConvNormAct(self.fea_dim[2], self.fea_dim[1], kernel_size=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),

        )
        self.up_4_2 = nn.Sequential(
            ConvNormAct(self.fea_dim[3], self.fea_dim[1], kernel_size=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),

        )
        self.up_4_3 = nn.Sequential(
            ConvNormAct(self.fea_dim[3], self.fea_dim[2], kernel_size=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),

        )

        self.conv1 = nn.Sequential(
            ConvNormAct(256, 64, kernel_size=1, stride=1),
            ConvNormAct(64, 64, kernel_size=3, stride=1, padding=1),
        )
        self.conv2 = nn.Sequential(
            ConvNormAct(448, 64, kernel_size=1, stride=1),
            ConvNormAct(64, 64, kernel_size=3, stride=1, padding=1),
        )
        self.conv3 = nn.Sequential(
            ConvNormAct(704, 128, kernel_size=1, stride=1),
            ConvNormAct(128, 128, kernel_size=3, stride=1, padding=1),
        )
        self.conv4 = nn.Sequential(
            ConvNormAct(960, 128, kernel_size=1, stride=1),
            ConvNormAct(128, 128, kernel_size=3, stride=1, padding=1),
        )









        self.mask_pred =nn.Conv2d(self.inter_channel*3, self.num_class, kernel_size=1, padding=0, bias=True)


    def forward(self, x):
        x1, x2, x3, x4 = self.backend(x)
        x2_1 = self.up_2_1(x2)
        x3_1 = self.up_3_1(x3)
        x4_1 = self.up_4_1(x4)

        x1_2 = self.down_1_2(x1)
        x3_2 = self.up_3_2(x3)
        x4_2 = self.up_4_2(x4)

        x1_3 = self.down_1_3(x1)
        x2_3 = self.down_2_3(x2)
        x4_3 = self.up_4_3(x4)

        x1_4 = self.down_1_4(x1)
        x2_4 = self.down_2_4(x2)
        x3_4 = self.down_3_4(x3)

        x1 = torch.cat([x1, x2_1, x3_1, x4_1],dim=1)
        x2 = torch.cat([x1_2, x2, x3_2, x4_2], dim=1)
        x3 = torch.cat([x1_3, x2_3, x3, x4_3], dim=1)
        x4 = torch.cat([x1_4, x2_4, x3_4,x4], dim=1)

        x1 = self.conv1(x1)
        x2 = self.conv2(x2)
        x3 = self.conv3(x3)
        x4 = self.conv4(x4)

        x4 = F.interpolate(x4, scale_factor=8, mode='bilinear', align_corners=False)
        x3 = F.interpolate(x3, scale_factor=4, mode='bilinear', align_corners=False)
        x2 = F.interpolate(x2, scale_factor=2, mode='bilinear', align_corners=False)

        x1 = torch.cat([x4, x3, x2, x1], dim=1)

        x1 = self.mask_pred(x1)
        x1 = F.interpolate(x1, scale_factor=4, mode='bilinear', align_corners=False)
        if self.is_edge:
            return torch.argmax(x1, dim=1, keepdim=True)
        else:
            return x1










if __name__ == '__main__':
    import torchstat

    model = SegUNet(num_class=10, load_pretrained=False)
    model = model
    model.eval()
    xs = torch.randn(1, 3, 544, 960)
    print(xs.size())
    m = model(xs)
    print(m.size())
